import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from data_interpolation import interpolate_data
from data import data
# 定义LSTM模型
class TrafficLSTM(nn.Module):
    def __init__(self, input_size=1, hidden_size=64, output_size=1, num_layers=1):
        super(TrafficLSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)
    
    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])
        return out

# 自定义数据集
class TrafficDataset(Dataset):
    def __init__(self, data, seq_length=5):
        self.data = data
        self.seq_length = seq_length
    
    def __len__(self):
        return len(self.data) - self.seq_length
    
    def __getitem__(self, idx):
        x = self.data[idx:idx+self.seq_length]
        y = self.data[idx+self.seq_length]
        return torch.FloatTensor(x), torch.FloatTensor([y])

# 示例数据预处理和训练流程
def main():
    print(torch.cuda.is_available())
    # 示例数据 - 替换为你的7组车流量数据
    traffic_data_non = np.array(data)
    
    traffic_data = interpolate_data(traffic_data_non, factor=4)
    # 数据标准化
    mean = np.mean(traffic_data)
    std = np.std(traffic_data)
    normalized_data = (traffic_data - mean) / std
    
    # 创建数据集
    seq_length = 128  # 使用前3个时间点预测下一个
    dataset = TrafficDataset(normalized_data.flatten(), seq_length)
    
    # 划分训练集和测试集
    train_size = int(0.8 * len(dataset))
    test_size = len(dataset) - train_size
    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])
    
    # 数据加载器
    batch_size = 128
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    
    # 初始化模型
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'使用设备: {device}')
    
    model = TrafficLSTM(input_size=1, hidden_size=64, output_size=1).to(device)
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    
    # 训练模型
    num_epochs = 200
    for epoch in range(num_epochs):
        for batch_x, batch_y in train_loader:
            batch_x = batch_x.unsqueeze(-1).to(device)
            batch_y = batch_y.to(device)
            
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        if (epoch+1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
    # 测试模型
    model.eval()
    with torch.no_grad():
        test_loss = 0
        for batch_x, batch_y in test_loader:
            batch_x = batch_x.unsqueeze(-1).to(device)
            batch_y = batch_y.to(device)
            outputs = model(batch_x)
            test_loss += criterion(outputs, batch_y).item()
            
            # 打印预测结果和真实值对比
            print(f'预测值: {outputs.squeeze().tolist()}, 真实值: {batch_y.squeeze().tolist()}')
        
        print(f'Test Loss: {test_loss/len(test_loader):.4f}')

if __name__ == '__main__':
    main()